package pogvue.analysis;

import pogvue.io.*;
import pogvue.datamodel.*;

import java.util.*;
import java.io.*;
import pogvue.gui.schemes.*;

public class Correlation2 {


    public static double get(double x[], double y[]) {

	double sum_sq_x = 0.0;
	double sum_sq_y = 0.0;
	double sum_coproduct = 0.0;
	double mean_x = x[2];
	double mean_y = y[2];

        int len = x.length;

	if (y.length < x.length) {
	  len = y.length;
	  }

        len -= 2;
	for (int i = 3; i < len; i++) {
	    double sweep = (i - 1.0) / i;
	    double delta_x = x[i] - mean_x;
	    double delta_y = y[i] - mean_y;
	    sum_sq_x += delta_x * delta_x * sweep;
	    sum_sq_y += delta_y * delta_y * sweep;
	    sum_coproduct += delta_x * delta_y * sweep;
	    mean_x += delta_x / i;
	    mean_y += delta_y / i;
	}
	double pop_sd_x = Math.sqrt( sum_sq_x / len);
	double pop_sd_y = Math.sqrt( sum_sq_y / len );
	double cov_x_y = sum_coproduct / len;
	return cov_x_y / (pop_sd_x * pop_sd_y);
    }

    public static double[] shift(double[] seq, int num) {
	double[] newseq = new double[seq.length + num*4];

	int i = 0;

	while (i < num) {
	    newseq[i*4]   = 0.25;
	    newseq[i*4+1] = 0.25;
	    newseq[i*4+2] = 0.25;
	    newseq[i*4+3] = 0.25;
	    i++;
	}
	i = i * 4;
	while (i < seq.length+num) {
	    newseq[i] = seq[i-num*4];
	    i++;
	}
	return newseq;
    }

    public static Hashtable get_with_shift(double[] vec1, double[] vec2,String seq,double thresh) {
	double[] off1 = Correlation2.shift(vec1,1);
	double[] off2 = Correlation2.shift(vec1,2);
	double[] off3 = Correlation2.shift(vec2,1);
	double[] off4 = Correlation2.shift(vec2,2);
	
	double corr =  Correlation2.get(vec1,vec2);
	double corr1 = Correlation2.get(vec2,off1);
	double corr2 = Correlation2.get(vec2,off1);
	double corr3 = Correlation2.get(vec1,off3);
	double corr4 = Correlation2.get(vec1,off4);
		    

	double maxcorr = corr;
	int    maxstr  = 0;
	
	if (corr1 > corr) {
	    if (corr2 > corr1) {
		maxstr = 2;
		maxcorr = corr2;
	    } else {
		maxstr = 1;
		maxcorr = corr1;
	    }
	}
	
	if (corr3 > maxcorr) {
	    if (corr4 > corr3) {
		maxstr = 4;
		maxcorr = corr4;
	    } else {
		maxstr = 3;
		maxcorr = corr3;
	    }
	}
	
	if (maxstr == 1) {
	    seq = "-" + seq; 
	} else if (maxstr == 2) {
	    seq = "--" + seq;
	} else if (maxstr == 3) {
	    seq = seq.substring(1) + "-";
	} else if (maxstr == 0) {
	//    seq = "--" + seq + "--";
	} else if (maxstr == 4) {
	    seq = seq.substring(2) + "--";
	}
	Hashtable out = new Hashtable();

	out.put("maxcorr",new Double(maxcorr));
	out.put("seq",seq);
        System.out.println("Maxcorr for " + seq + " " + maxcorr); 
	return out;
    }

    public static void main(String[] args) {

	try {
	    FastaFile ff     = null;
	    PwmFile   pf     = null;
	    Vector    seqs   = null;
	    Vector    pwm    = null;

	    double    thresh = Double.valueOf(args[1]);
	    String    type   = args[2];


	    if (type.equals("-fasta")) {
		ff = new FastaFile(args[0],"File");
		seqs   = ff.getSeqs();
	    } else if (type.equals("-pwm")) {
		pf = new PwmFile(args[0],"File",12);
		pwm  = pf.getPwm();
	    }
		
	    Vector    clus   = new Vector();

	    for (int i = 0; i < seqs.size(); i++) {

		String seq = ((Sequence)seqs.elementAt(i)).getSequence();		
		seq = "--" + seq + "--";

		double[] vec1 = Correlation2.seqvec(new Sequence(seq,seq,1,seq.length()));

		System.out.print("Sequence " + seq + "\n");
		//Correlation2.printPWM(vec1);			
		//Pwm.printLogo(vec1);
		//System.out.println();

		int found = 0;

		for (int j = 0; j < clus.size(); j++) {
		    pogvue.datamodel.Cluster c = (pogvue.datamodel.Cluster)clus.elementAt(j);
		    Vector clusseqs = (Vector)c.getSeqs();
                    System.out.println("Cons " + c.getConsensus()); 
		    double[] vec2 = c.getPwm();
		    //  System.out.print("Pwm ");
		    //Correlation2.printPWM(vec2);
		    //Pwm.printLogo(vec2);
		    //System.out.println();

		    Hashtable res = Correlation2.get_with_shift(vec1,vec2,seq,thresh);

		    double maxcorr = ((Double)res.get("maxcorr")).doubleValue();
		    String newseq  = (String)res.get("seq");

		    String revstr = ResidueProperties.reverseComplement(seq);

		    double[] revvec1 = Correlation2.seqvec(new Sequence(revstr,revstr,1,revstr.length()));

		    Hashtable revres = Correlation2.get_with_shift(revvec1,vec2,revstr,thresh);

		    
		    double revmaxcorr = ((Double)revres.get("maxcorr")).doubleValue();
		    String revnewseq  = (String)revres.get("seq");

		    if (revmaxcorr > maxcorr) {
			maxcorr = revmaxcorr;
			newseq  = revnewseq;
		    }

		    if (maxcorr > thresh) {
			System.out.println("Add to cluster " + newseq + "\t" + c.getConsensus() + "\t" + maxcorr + "\t" + clusseqs.size());
			found = 1;
			c.addSequence(new Sequence(newseq,newseq,1,newseq.length()));
			System.out.println("Updated pwm");

			Correlation2.printPWM(c.getPwm(),false);
			Pwm.printLogo(c.getPwm());

			System.out.println();
			j = clus.size();
		    }
		}

		if (found == 0) {
		    String tmpseq = ((Sequence)seqs.elementAt(i)).getSequence();
		    tmpseq = "--" + tmpseq + "--";
		    pogvue.datamodel.Cluster newclus = new pogvue.datamodel.Cluster(new Sequence(tmpseq,tmpseq,1,tmpseq.length()));
		    clus.addElement(newclus);
		} 

	    }

	    for (int i = 0; i < clus.size(); i++) {
		pogvue.datamodel.Cluster c = (pogvue.datamodel.Cluster)clus.elementAt(i);
		Vector  cseqs = c.getSeqs();

		System.out.print("Cluster\t" + cseqs.size());

		for (int j = 0; j < cseqs.size(); j++) {
			System.out.print("\t" + ((Sequence)cseqs.elementAt(j)).getSequence());
		}
		System.out.print("\t");
		Correlation2.printPWM(c.getPwm(),true);
		System.out.println();

		Pwm.printLogo(c.getPwm());
		System.out.println();
		System.out.println(c.getConsensus());
		System.out.println();
	    }
		    
	} catch (IOException e) {
	    System.out.println("Exception " + e);
	}

    }
    
    public static double[] seqvec(Sequence seq) {
	
	double seqvec[] = new double[seq.getLength()*4];
	
	for (int j = 0; j < seq.getLength(); j++) {
	    char c = seq.getCharAt(j);
	    
	    seqvec[j*4] = 0;
	    seqvec[j*4+1] = 0;
	    seqvec[j*4+2] = 0;
	    seqvec[j*4+3] = 0;
	    if (c == 'A') {
		seqvec[j*4] = 1;
	    } else if (c == 'T') {
		seqvec[j*4+1] = 1;
	    } else if (c == 'C') {
		seqvec[j*4+2] = 1;
	    } else if (c == 'G') {
		seqvec[j*4+3] = 1;
	    } else if (c == '-') {
	        seqvec[j*4] = 0.25;
		seqvec[j*4+1] = 0.25;
		seqvec[j*4+2] = 0.25;
		seqvec[j*4+3] = 0.25;
		}
	}
	return seqvec;
    }
    public static String vec2seq(double[] vec) {

	StringBuffer str = new StringBuffer();

	int i = 0;
	
	while (i < vec.length/4) {
	    int max = maxOfFour(vec,i);

	    if (i == 0) {
		str.append("A");
	    } else if (i == 1) {
		str.append("T");
	    } else if (i == 2) {
		str.append("C");
	    } else if (i ==1) {
		str.append("G");
	    }
	    i += 4;
	}
	return str.toString();
    }

    public static int maxOfFour(double[] vec, int offset) {

	double max = 0;
	int maxindex = -1;

	int i = 0;

	while (i < 4) {
	    if (vec[offset+i] >= max) {
		max = vec[offset+i];
		maxindex = i;
	    }
	    i++;
	}
	return maxindex;
    }
    public static void printPWM(double[] pwm, boolean oneline) {

	int i = 0;

	while (i < pwm.length/4) {
	    int j = 0;

	    while (j < 4) {
		System.out.printf("%10.2f\t",pwm[i*4+j]);
		j++;
	    }
	    if (oneline == false) {
		System.out.println();
	    }
	    i++;
	}
	
    }
}
